#include "aim_rec/NumRec.hpp"
#include "aim_deps/Paths.hpp"
#include "readParam.hpp"
NumRec::NumRec()
{
    std::string stn_path = model_path + "stn.pt";
    stn = torch::jit::load(stn_path);
}

void NumRec::reco(aim_deps::Armor &Input_armor, const cv::Mat &grayImg, int type)
{
    cv::Point2f pts[4];
    getExtendedVex(Input_armor.left_light.box, Input_armor.right_light.box, pts);
    const cv::Point2f &
    tl = pts[0],
    bl = pts[1],
    br = pts[2],
    tr = pts[3];

    int width, height;
    
    width = 28;
    height = 28;
    
    cv::Point2f src[4]{cv::Vec2f(tl), cv::Vec2f(tr), cv::Vec2f(br), cv::Vec2f(bl)};                                           //实际矩形
    cv::Point2f dst[4]{cv::Point2f(0.0, 0.0), cv::Point2f(width, 0.0), cv::Point2f(width, height), cv::Point2f(0.0, height)}; //目标矩形
    const cv::Mat perspMat = cv::getPerspectiveTransform(src, dst);                                                           //变换
    cv::Mat _transformed_Img;                                                                                                 //the Image of the aim_deps::Armor
    cv::warpPerspective(grayImg, _transformed_Img, perspMat, cv::Size(width, height));                                        //获取矫正图像    

    // cv::imshow("_transformed_Img", _transformed_Img);
    // cv::threshold(_transformed_Img, _transformed_Img, 50, 255, cv::THRESH_BINARY);                                            //阈值化，准备识别

    //数据增强
    // for(int i = 0; i < height ; i++){
    //     for(int j = 0; j < width; j++){
    //         _transformed_Img.at<uchar>(i,j) = _transformed_Img.at<uchar>(i,j)+2;
    //     }
    // }

    cv::threshold(_transformed_Img, _transformed_Img, light_theshold, 255, cv::THRESH_BINARY);
    // //如果要输出结果
    // if (Input_armor.Isbigarmor){
    //     std::sprintf(str, "/home/xjturm/cnn/cnn_big_new/1/img%d.jpg", ++img_cnt);
    //     imwrite(str, _transformed_Img);
    //     std::cout<<width<<std::endl;
    // }
    

    //不同阈值化的尝试
    // cv::threshold(_transformed_Img, _transformed_Img, 40, 255, cv::THRESH_TOZERO_INV);
    // // cv::equalizeHist(_transformed_Img, _transformed_Img);
    // // cv::imshow("equalizeHist", _transformed_Img);
    // cv::threshold(_transformed_Img, _transformed_Img, 2, 255, cv::THRESH_TOZERO);
    //  cv::threshold(_transformed_Img, _transformed_Img, 0, 255, cv::THRESH_BINARY | cv::THRESH_OTSU);
    

    cv::resize(_transformed_Img, _transformed_Img, cv::Size(width, height), 0, 0, cv::INTER_AREA); //防止可能的bug
    
    
    // cv::waitKey(0);

    int r = std::round(Classfier(_transformed_Img,Input_armor.Isbigarmor)) + 1;
    // std::cout << r << std::endl;
    //大小装甲板分类识别
    // if(Input_armor.Isbigarmor){
    //     switch (r)
    //     {
    //     case 0:
    //         Input_armor.armor_number = 1;//英雄
    //         break;
    //     case 1:
    //         Input_armor.armor_number = 7;//哨兵
    //         break;
    //     case 2:
    //         Input_armor.armor_number = 8;//基地
    //         break;
    //     case 3:
    //         Input_armor.armor_number = -1;//什么都不是
    //         break;
    //     default:
    //         break;
    //     }
    // }
    // else{
    //     switch (r)
    //     {
    //     case 0:
    //         Input_armor.armor_number = -1;//什么都不是
    //         break;
    //     case 1:
    //         Input_armor.armor_number = 2;//工程
    //         break;
    //     case 2:
    //         Input_armor.armor_number = 3;//步兵3
    //         break;
    //     case 3:
    //         Input_armor.armor_number = 4;//步兵4
    //         break;
    //     case 4:
    //         Input_armor.armor_number = 5;//步兵5
    //         break;
    //     case 5:
    //         Input_armor.armor_number = 6;//基地小装甲
    //         break;
    //     }
    // }
    Input_armor.armor_number = r;
    // cv::imshow("a", _transformed_Img);

}

void NumRec::getExtendedVex(
    aim_deps::LightBox l1,
    aim_deps::LightBox l2,
    cv::Point2f pts[]) const
{
    const int x_shorten_ratio =0.3;
    l1.extend(2.0);
    l2.extend(2.0);
    // pts[0] = cv::Point2f(l1.vex[0].x + (l2.vex[0].x - l1.vex[0].x) * x_shorten_ratio, l1.vex[0].y);
    // pts[1] = cv::Point2f(l1.vex[1].x + (l2.vex[1].x - l1.vex[1].x) * x_shorten_ratio, l1.vex[1].y);
    // pts[2] = cv::Point2f(l2.vex[1].x - (l2.vex[1].x - l1.vex[1].x) * x_shorten_ratio, l2.vex[1].y);
    // pts[3] = cv::Point2f(l2.vex[0].x - (l2.vex[0].x - l1.vex[0].x) * x_shorten_ratio, l2.vex[0].y);
    pts[0] = l1.vex[0];
    pts[1] = l1.vex[1];
    pts[2] = l2.vex[1];
    pts[3] = l2.vex[0];
}

float NumRec::Classfier(cv::Mat &image,bool isbig)
{
    torch::Tensor img_tensor = torch::from_blob(image.data, {1, image.rows, image.cols, 1}, torch::kByte);
    img_tensor = img_tensor.permute({0, 3, 1, 2});
    img_tensor = img_tensor.toType(torch::kFloat);
    img_tensor = img_tensor.div(255);
    torch::Tensor output;
    // if(isbig){
    //     output = cnn_big.forward({img_tensor}).toTensor();
    // }else{
    //     output = cnn_small.forward({img_tensor}).toTensor();
    // }
    output = stn.forward({img_tensor}).toTensor();
    auto max_result = output.max(1, true);
    return std::get<1>(max_result).item<float>();
}